import json
import logging
import pandas as pd
import re
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

import os

from catalog.dip_catalog import DipCatalog
from util.read_txt_file import ReadTxt
from series.agent.age_supplementary_catalog import AgeSupplementaryCatalog

logger = logging.getLogger(__name__)

class DIPGroupRules:
    def __init__(self):
        logger.info("开始初始化大模型")

        # 从环境变量读取配置
        load_dotenv()
        api_key = os.getenv("OPENAI_API_KEY")
        base_url = os.getenv("OPENAI_BASE_URL")
        model_name = os.getenv("MODEL_NAME")
        logger.info("初始化DIP分组规则系统")
        
        # 初始化年龄智能体
        self.age_agent = AgeSupplementaryCatalog()
        
        # 初始化大模型
        self.llm = ChatOpenAI(
            temperature=0.1,
            max_tokens=300,
            base_url=base_url,
            model=model_name,
            api_key=api_key
        )
        logger.info("大模型初始化完成")
        self.query_prompt = ChatPromptTemplate.from_template(
            """你是DIP病种分组专家，请根据DIP病种目录库数据进行精确分组。

患者信息：
- 诊断编码：{diagnosis_code}
- 诊断名称：{diagnosis_name}
- 患者年龄：{patient_age}岁
- 住院天数：{hospitalization_days}天
- 主手术编码：{main_surgery_code}
- 次要手术编码：{other_surgery_code}
- 治疗操作：{treatment_code}
- 治疗操作：{diagnostic_operation}

DIP目录库匹配数据：
{catalog_data}

统一入组规则：
{unified_rules}

重要规则：
1. 如果没有手术编码,没有治疗编码,没有诊断操作编码，应返回"核心保守治疗"
2. 如果有手术编码，应返回对应手术的治疗方式名称,如果治疗操作编码不为空,返回对应治疗操作相关
3. DIP编码格式为"诊断编码:手术编码"，如A09.9:45.2501
4. **重要**：请从目录库匹配数据中选择最合适的记录，并使用该记录中的"分值"列作为分值结果
5. 床日分析：根据患者住院天数判断是否需要床日辅助调整（通常住院天数大于21才可以）
6. 年龄分析：根据患者年龄判断是否需要年龄辅助调整

**分值获取规则**：
- 优先匹配患者的手术编码或治疗操作编码对应的DIP记录
- 如果有精确匹配的DIP编码，使用该记录的分值
- 如果没有精确匹配，选择诊断编码下最合适的治疗方式记录
- 分值必须从目录库数据中获取，不能为空

请根据目录库数据，返回以下格式的结果：
分组名称：[具体的治疗方式名称或"核心保守治疗"]
分值：[数值]
手术及操作编码：[手术及操作编码]
诊断编码:[数值]

需要床日调整：[是/否]
床日调整原因：[如果需要调整，说明原因；如果不需要则写"无"]

需要年龄调整：[是/否]
年龄调整原因：[患者年龄调整范围:0-28天,29天-1岁,1-6岁,7-17岁,66-120岁]

**注意**：如果判断需要床日或年龄调整，只需要标明"是"和原因，具体的分值调整将由后续的专门模块处理。

结果："""
        )
        # 创建处理链
        self.chain = self.query_prompt | self.llm | StrOutputParser()

    def parse_ai_result(self, ai_result_text):
        """
        解析AI返回的文本格式，提取结构化数据
        
        Args:
            ai_result_text (str): AI返回的原始文本
            
        Returns:
            dict: 解析后的结构化数据
        """
        try:
            result_data = {
                "分组名称": "",
                "分值": None,
                "手术及操作编码": "",
                "诊断编码": "",
                "需要床日调整": False,
                "床日调整原因": "",
                "需要年龄调整": False,
                "年龄调整原因": "",
                "原始结果": ai_result_text
            }
            
            # 使用正则表达式提取各个字段
            patterns = {
                "分组名称": r"分组名称：(.+?)(?:\n|$)",
                "分值": r"分值：(.+?)(?:\n|$)",
                "手术及操作编码": r"手术及操作编码：(.+?)(?:\n|$)",
                "诊断编码": r"诊断编码：(.+?)(?:\n|$)",
                "需要床日调整": r"需要床日调整：(.+?)(?:\n|$)",
                "床日调整原因": r"床日调整原因：(.+?)(?:\n|$)",
                "需要年龄调整": r"需要年龄调整：(.+?)(?:\n|$)",
                "年龄调整原因": r"年龄调整原因：(.+?)(?:\n|$)"
            }
            
            for field, pattern in patterns.items():
                match = re.search(pattern, ai_result_text, re.IGNORECASE)
                if match:
                    value = match.group(1).strip()
                    
                    if field == "分值":
                        try:
                            result_data[field] = float(value)
                        except ValueError:
                            logger.warning(f"无法转换分值为数字: {value}")
                            result_data[field] = None
                    elif field in ["需要床日调整", "需要年龄调整"]:
                        result_data[field] = value.lower() in ["是", "yes", "true"]
                    else:
                        result_data[field] = value
            
            logger.info(f"成功解析AI结果: {result_data}")
            return result_data
            
        except Exception as e:
            logger.error(f"解析AI结果时发生错误: {str(e)}")
            return {
                "分组名称": "解析失败",
                "分值": None,
                "手术及操作编码": "",
                "诊断编码": "",
                "需要床日调整": False,
                "床日调整原因": "",
                "需要年龄调整": False,
                "年龄调整原因": "",
                "原始结果": ai_result_text,
                "解析错误": str(e)
            }

    def get_surgery_code(self, patient_row, catalog_data):
        """
        便捷方法：直接获取手术及操作编码
        
        Args:
            patient_row: 患者数据
            catalog_data: 目录数据
            
        Returns:
            str: 手术及操作编码
        """
        try:
            ai_result = self.execute_ai(patient_row, catalog_data)
            return ai_result['结构化数据'].get('手术及操作编码', '')
        except Exception as e:
            logger.error(f"获取手术编码失败: {str(e)}")
            return ""
    
    def get_group_name(self, patient_row, catalog_data):
        """
        便捷方法：直接获取分组名称
        
        Args:
            patient_row: 患者数据
            catalog_data: 目录数据
            
        Returns:
            str: 分组名称
        """
        try:
            ai_result = self.execute_ai(patient_row, catalog_data)
            return ai_result['结构化数据'].get('分组名称', '')
        except Exception as e:
            logger.error(f"获取分组名称失败: {str(e)}")
            return ""
    
    def get_score(self, patient_row, catalog_data):
        """
        便捷方法：直接获取分值
        
        Args:
            patient_row: 患者数据
            catalog_data: 目录数据
            
        Returns:
            float: 分值
        """
        try:
            ai_result = self.execute_ai(patient_row, catalog_data)
            return ai_result['结构化数据'].get('分值')
        except Exception as e:
            logger.error(f"获取分值失败: {str(e)}")
            return None
    
    def need_age_adjustment(self, patient_row, catalog_data):
        """
        便捷方法：判断是否需要年龄调整
        
        Args:
            patient_row: 患者数据
            catalog_data: 目录数据
            
        Returns:
            bool: 是否需要年龄调整
        """
        try:
            ai_result = self.execute_ai(patient_row, catalog_data)
            return ai_result['结构化数据'].get('需要年龄调整', False)
        except Exception as e:
            logger.error(f"判断年龄调整需求失败: {str(e)}")
            return False
    
    def need_bed_day_adjustment(self, patient_row, catalog_data):
        """
        便捷方法：判断是否需要床日调整
        
        Args:
            patient_row: 患者数据
            catalog_data: 目录数据
            
        Returns:
            bool: 是否需要床日调整
        """
        try:
            ai_result = self.execute_ai(patient_row, catalog_data)
            return ai_result['结构化数据'].get('需要床日调整', False)
        except Exception as e:
            logger.error(f"判断床日调整需求失败: {str(e)}")
            return False
    
    def extract_field(self, ai_result, field_name):
        """
        从AI结果中提取特定字段的通用方法
        
        Args:
            ai_result: execute_ai返回的结果
            field_name: 字段名称
            
        Returns:
            提取的字段值
        """
        if isinstance(ai_result, dict) and '结构化数据' in ai_result:
            return ai_result['结构化数据'].get(field_name)
        else:
            logger.warning("AI结果格式不正确")
            return None

    # 执行ai
    def execute_ai(self,patient_row,catalog_data):
        unified_rules = self.read_knowledge_base_txt()
        try:
            res = json.loads(catalog_data)
            first_record = res["data"][0]
            
            # 执行基本的DIP分组
            result = self.chain.invoke({
                "patient_info": patient_row,
                "diagnosis_code": patient_row['诊断编码'],
                "diagnosis_name": first_record['诊断名称'],
                'hospitalization_days': patient_row['住院天数'],
                "patient_age": patient_row.get('年龄', 0),
                "main_surgery_code":patient_row['主手术'] if pd.notna(patient_row['主手术']) else "",
                "other_surgery_code": patient_row['其它手术'] if pd.notna(patient_row['其它手术']) else "",
                "treatment_code": patient_row['治疗操作'] if pd.notna(patient_row['治疗操作']) else "",
                "diagnostic_operation": patient_row['诊断操作'] if pd.notna(patient_row['诊断操作']) else "",
                "catalog_data": catalog_data,
                "unified_rules": unified_rules[:1200]
            })
            logger.info("AI大模型分析完成")
            logger.info(f"AI返回结果: {result}")
            


            # 解析AI结果为结构化数据
            parsed_result = self.parse_ai_result(result)

            if self.ai_check_need_age_agent(parsed_result):
                patient_age = patient_row.get('年龄', 0)
                logger.info(f"AI判断需要调用年龄智能体，患者年龄{patient_age}岁")
                result = self._call_age_agent(patient_age, parsed_result)

            
            # 返回结构化数据和原始文本
            return {
                "原始文本": result.strip(),
                "结构化数据": parsed_result
            }



        except Exception as e:
            logger.error(f"执行AI时发生错误: {str(e)}")
            logger.error(f"错误类型: {type(e).__name__}")
            print(f"执行AI时发生错误: {str(e)}")
            print(f"错误类型: {type(e).__name__}")
            return {
                "原始文本": "",
                "结构化数据": {
                    "分组名称": "执行失败",
                    "分值": None,
                    "错误信息": str(e)
                }
            }
    
    def ai_check_need_age_agent(self, ai_result):

        dipCatalog =DipCatalog()
        surgery_code = ai_result.get('手术及操作编码', '')
        print("*******************",surgery_code)

        try:
            dipCatalog.load_age_supplementary_catalog()
            # 检查AI返回的结果中是否包含需要年龄调整的标识
            if ai_result.get('需要年龄调整', False):
                return True

            return False
        except Exception as e:
            logger.error(f"AI判断年龄智能体需求失败: {str(e)}")
            return False
    
    def _call_age_agent(self, patient_age, base_result):
        """调用年龄智能体"""
        try:
            logger.info(f"调用年龄智能体，患者年龄: {patient_age}岁")
            
            # 调用简化的年龄智能体
            age_result = self.age_agent.analyze_age_adjustment(patient_age, base_result)
            
            # 将年龄调整结果添加到原结果中
            return base_result + f"\n\n【年龄智能体调整】\n{age_result}"
            
        except Exception as e:
            logger.error(f"年龄智能体调用失败: {str(e)}")
            return base_result

    # 读取知识库(txt)
    def read_knowledge_base_txt(self):
        reader_txt = ReadTxt('规则文件.txt')
        return reader_txt.load_txt_lines()